#include<stdio.h>
#include<stdlib.h>
#include"fftz.h"
#include<clAmdFft.h>

static int verifyAmd(FftzEnv* cl, int N, float* data0){
	clAmdFftSetupData fft_setup_data;
	clAmdFftPlanHandle fft_plan;
	int err;
	size_t size = N;
	
	cl_mem buffer0[2]; 
	cl_mem buffer1[2];
	
	buffer0[0] = clCreateBuffer(cl->context, CL_MEM_READ_WRITE, sizeof(float) * N, NULL, &err);
	if(err < 0){
		printf("Error: Failed to create buffer0[0]!\n");
		return -1;
	}
	buffer0[1] = clCreateBuffer(cl->context, CL_MEM_READ_WRITE, sizeof(float) * N, NULL, &err);
	if(err < 0){
		printf("Error: Failed to create buffer0[1]!\n");
		clReleaseMemObject(buffer0[0]);
		return -1;
	}
	buffer1[0] = clCreateBuffer(cl->context, CL_MEM_READ_WRITE, sizeof(float) * N, NULL, &err);
	if(err < 0){
		printf("Error: Failed to create buffer1[0]!\n");
		clReleaseMemObject(buffer0[0]);
		clReleaseMemObject(buffer0[1]);
		return -1;
	}
	buffer1[1] = clCreateBuffer(cl->context, CL_MEM_READ_WRITE, sizeof(float) * N, NULL, &err);
	if(err < 0){
		printf("Error: Failed to create buffer1[1]!\n");
		clReleaseMemObject(buffer0[0]);
		clReleaseMemObject(buffer0[1]);
		clReleaseMemObject(buffer1[0]);
		return -1;
	}
	
	err = clEnqueueWriteBuffer(cl->queue, buffer0[0], CL_TRUE, 0, sizeof(float) * N, data0, 0, NULL, NULL);
	if(err < 0){
		printf("Error: Failed to Write buffer0!\n");
		clReleaseMemObject(buffer0[0]);
		clReleaseMemObject(buffer0[1]);
		clReleaseMemObject(buffer1[0]);
		clReleaseMemObject(buffer1[1]);
		return -1;
	}
	err = clEnqueueWriteBuffer(cl->queue, buffer0[1], CL_TRUE, 0, sizeof(float) * N, data0 + N, 0, NULL, NULL);
	if(err < 0){
		printf("Error: Failed to Write buffer0!\n");
		clReleaseMemObject(buffer0[0]);
		clReleaseMemObject(buffer0[1]);
		clReleaseMemObject(buffer1[0]);
		clReleaseMemObject(buffer1[1]);
		return -1;
	}
	
	// Get FFT version  
	err = clAmdFftInitSetupData(&fft_setup_data);
	if(err < 0){
		printf("Error: Failed to setup amdfft!\n");
		clReleaseMemObject(buffer0[0]);
		clReleaseMemObject(buffer0[1]);
		clReleaseMemObject(buffer1[0]);
		clReleaseMemObject(buffer1[1]);
		return -1;
	}
	//printf("Using clAmdFft %u.%u.%u\n",fft_setup_data.major,fft_setup_data.minor,fft_setup_data.patch);

	//FFT Setup
	err = clAmdFftSetup(&fft_setup_data);
	if(err < 0){
		printf("Error: Failed to setup amdfft!\n");
		clReleaseMemObject(buffer0[0]);
		clReleaseMemObject(buffer0[1]);
		clReleaseMemObject(buffer1[0]);
		clReleaseMemObject(buffer1[1]);
		return -1;
	}

	// Create FFT plan  
	err = clAmdFftCreateDefaultPlan(&fft_plan, cl->context, CLFFT_1D, &size);
	if(err < 0){
		printf("Error: Failed to create default plan!\n");
		clAmdFftTeardown();
		clReleaseMemObject(buffer0[0]);
		clReleaseMemObject(buffer0[1]);
		clReleaseMemObject(buffer1[0]);
		clReleaseMemObject(buffer1[1]);
		return -1;
	}
	
	clAmdFftSetPlanPrecision(fft_plan, CLFFT_SINGLE);
	clAmdFftSetPlanBatchSize(fft_plan, 1);
	clAmdFftSetLayout(fft_plan, CLFFT_COMPLEX_PLANAR, CLFFT_COMPLEX_PLANAR);
	clAmdFftSetResultLocation(fft_plan, CLFFT_OUTOFPLACE);
	
	err = clAmdFftEnqueueTransform(fft_plan, CLFFT_FORWARD, 1, &cl->queue, 0, NULL, NULL, buffer0, buffer1, NULL);
	if(err < 0){
		printf("Error: Failed to enqueue transform!\n");
		clAmdFftDestroyPlan(&fft_plan);
		clAmdFftTeardown();
		clReleaseMemObject(buffer0[0]);
		clReleaseMemObject(buffer0[1]);
		clReleaseMemObject(buffer1[0]);
		clReleaseMemObject(buffer1[1]);
		return -1;
	}
	clFinish(cl->queue);
	
	err = clEnqueueReadBuffer(cl->queue, buffer1[0], CL_TRUE, 0, sizeof(float) * N, data0, 0, NULL, NULL);
	if(err < 0){
		printf("Error: Failed to read buffer1!\n");
		clAmdFftDestroyPlan(&fft_plan);
		clAmdFftTeardown();
		clReleaseMemObject(buffer0[0]);
		clReleaseMemObject(buffer0[1]);
		clReleaseMemObject(buffer1[0]);
		clReleaseMemObject(buffer1[1]);
		return -1;
	}
	err = clEnqueueReadBuffer(cl->queue, buffer1[1], CL_TRUE, 0, sizeof(float) * N, data0 + N, 0, NULL, NULL);
	if(err < 0){
		printf("Error: Failed to read buffer1!\n");
		clAmdFftDestroyPlan(&fft_plan);
		clAmdFftTeardown();
		clReleaseMemObject(buffer0[0]);
		clReleaseMemObject(buffer0[1]);
		clReleaseMemObject(buffer1[0]);
		clReleaseMemObject(buffer1[1]);
		return -1;
	}
	clFinish(cl->queue);
		
	clAmdFftDestroyPlan(&fft_plan);
	clAmdFftTeardown();	
	clReleaseMemObject(buffer0[0]);
	clReleaseMemObject(buffer0[1]);
	clReleaseMemObject(buffer1[0]);
	clReleaseMemObject(buffer1[1]);
	
	return 0;
}

int main(int argc, char** argv){
	FftzEnv cl;
	FftzPlan plan;
	float *data0;
	int i;
	FILE* file;
	int N;
	char outFile[100];
	int err;

	if(initFftzEnv(&cl) < 0){
		printf("Initialization of FftzEnv failed, program exiting..\n");
		return -1;
	}

	if(argc != 3){
		printf("Argument count doesn't match, program exiting..\n");
		releaseFftzEnv(&cl);
		return -1;
	}
	
	N = 1 << (atoi(argv[2]));
	data0 = (float*)malloc(sizeof(float) * N * 2);
	printf("FFT size selected: %d\n", N);

	for(i = 0; i < N; i++){
		data0[i] = i * i;
		data0[i + N] = i;
	}
	
	switch(argv[1][0]){
		case 'z':
			printf("FFT function selected: fftz\n");
			sprintf(outFile, "output/verifyFftz%d", N);
			err = initFftzPlan(&cl, &plan, atoi(argv[2]));
			if(err < 0){
				printf("Initialization of FftzPlan failed with error code %d, program exitting..\n", err);
				free(data0);
				releaseFftzEnv(&cl);
				return -1;
			}
			err = fftzRun(&cl, &plan, data0);
			if(err < 0){
				printf("function fftzRun failed with error code %d, program exitting..\n", err);
				releaseFftzPlan(&plan);
				free(data0);
				releaseFftzEnv(&cl);
				return err;
			}
			releaseFftzPlan(&plan);
			break;
		case 'a':
			printf("FFT algorithm selected: AmdClFft\n");
			sprintf(outFile, "output/verifyAmd%d", N);
			if(verifyAmd(&cl, N, data0) < 0){
				printf("function verifyAmd returned error, program exiting..\n");
				free(data0);
				releaseFftzEnv(&cl);
				return -1;
			}
			break;
		default:
			printf("Cannot recognize the first argument, program exiting..\n");
			free(data0);
			releaseFftzEnv(&cl);
			return -1;
	}
	
	
	if(N <= 131072){
		printf("Caculation completed.\nWriting results to %s.\n", outFile);
		file = fopen(outFile, "w");
		if(file == NULL){
			printf("Failed to open %s, program exiting..\n", outFile);
			free(data0);
			releaseFftzEnv(&cl);
			return -1;
		}
		for(i = 0; i < N; i++){
			fprintf(file, "%.4f %.4f\n", data0[i], data0[i + N]);
		}	
		fclose(file);
	}
	else{
		printf("FFT size %d too large, choose not to write results to disk.", N);
	}
	
	free(data0);
	releaseFftzEnv(&cl);
	
	return 0;
}
